import torch
from torchvision import transforms
from PIL import Image
import sys
import os
import json

# 添加项目根目录到系统路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from models.simple_cnn import SimpleCNN

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 加载类别信息和建议
config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'class_info.json')
with open(config_path, 'r', encoding='utf-8') as f:
    class_info = json.load(f)
class_names = class_info['class_names']
suggestions = class_info['suggestions']

# 定义模型结构
num_classes = len(class_names)
model = SimpleCNN(num_classes).to(device)

# 加载模型检查点
checkpoint_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'checkpoint_epoch_25.pth.tar')
checkpoint = torch.load(checkpoint_path, map_location=device)
model.load_state_dict(checkpoint['state_dict'])

# 将模型设置为评估模式
model.eval()

# 定义数据转换
test_transforms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 处理图片并获取预测结果
def predict_image(image_path, top_k=3):
    try:
        image = Image.open(image_path).convert('RGB')  # 确保图像为RGB格式
        image = test_transforms(image).unsqueeze(0).to(device)  # 添加批次维度
    except Exception as e:
        print(f"Error processing image {image_path}: {e}")
        return [], "图片处理失败，请确认文件是否为有效的图片格式。"

    # 获取预测结果
    with torch.no_grad():
        output = model(image)
        probabilities = torch.nn.functional.softmax(output, dim=1)[0]
        
        # 获取 top-k 结果
        top_probs, top_indices = torch.topk(probabilities, top_k)

    results = []
    for i in range(top_k):
        prob = top_probs[i].item()
        class_idx = top_indices[i].item()
        class_name = class_names[class_idx]
        suggestion = suggestions.get(class_name, "暂无防治建议")
        results.append({
            "class_name": class_name,
            "probability": prob,
            "suggestion": suggestion
        })

    # 添加可信度检查逻辑
    additional_message = ""
    top_prediction_prob = results[0]['probability']
    if top_prediction_prob < 0.7:
        additional_message = "识别图片可能非植物叶片。"
    elif top_prediction_prob < 0.9:
        additional_message = "识别物体可能非苹果叶片。"

    return results, additional_message

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("使用方法: python predict.py <图片路径>")
        sys.exit(1)
    
    image_path = sys.argv[1]
    if not os.path.exists(image_path):
        print(f"错误：找不到图片文件 {image_path}")
        sys.exit(1)
    
    try:
        results, additional_message = predict_image(image_path)
        if not results:
            print(additional_message)
            sys.exit(1)

        print("\n预测结果：")
        for res in results:
            print(f"  - 类别: {res['class_name']}, 置信度: {res['probability']:.2%}")
            print(f"    建议: {res['suggestion']}")
        
        if additional_message:
            print(f"\n注意: {additional_message}")
    except Exception as e:
        print(f"预测过程中出现错误: {str(e)}")
        sys.exit(1)
